import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import jieba
import re

# 检查是否有可用的CUDA设备，将设备设置为GPU或CPU
from torch import Tensor

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def batchify(data: Tensor, bsz: int) -> Tensor:
    seq_len = data.size(0) // bsz
    data = data[:seq_len * bsz]
    data = data.view(bsz, seq_len)
    #data = data.view(bsz, seq_len).t().contiguous()
    return data.to(device)


def split_chinese(sentence):
    # 使用jieba进行分词
    words = jieba.lcut(sentence)
    # 利用正则表达式过滤汉字以外的字符
    pattern = re.compile(r'[\w\s]')
    words = [word for word in words if pattern.search(word)]
    return words


# 创建数据加载器
class TextDataset(torch.utils.data.Dataset):
    def __init__(self, data):
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx, :-1], self.data[idx, 1:]


# 定义位置编码
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super(PositionalEncoding, self).__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-torch.log(torch.tensor(10000.0)) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        return x + self.pe[:x.size(0), :]


# 定义 GPT 模型
class SimpleGPT(nn.Module):
    def __init__(self, vocab_size, d_model=256, nhead=8, num_decoder_layers=3):
        super(SimpleGPT, self).__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.positional_encoding = PositionalEncoding(d_model)
        decoder_layer = nn.TransformerDecoderLayer(d_model, nhead)
        self.transformer = nn.TransformerDecoder(decoder_layer, num_decoder_layers)
        self.fc = nn.Linear(d_model, vocab_size)

    def forward(self, src, tgt):
        src = self.embedding(src) * np.sqrt(256)
        src = self.positional_encoding(src)

        tgt = self.embedding(tgt) * np.sqrt(256)
        tgt = self.positional_encoding(tgt)
        # 将 None 传递给 self.transformer 的第二个参数是不正确的，因为 Transformer 模型需要目标序列作为输入
        memory = self.transformer(src, tgt)  # 使用tgt作为目标序列
        output = self.fc(memory)
        return output


# 测试问答系统
def generate_answer(question, model, vocab):
    # 将问题转换为模型可接受的输入格式
    question_words = split_chinese(question)
    input_data = [vocab[word] for word in question_words if word in vocab]
    input_tensor = torch.tensor(input_data).unsqueeze(0)  # 添加 batch 维度

    # 在这里假设我们的模型可以生成固定长度的输出
    target_length = 10  # 假设生成包含10个单词的回答

    model.eval()
    with torch.no_grad():
        output_sequence = []  # 存储模型生成的输出序列
        for i in range(target_length):
            output = model(input_tensor, input_tensor)  # 这里将输入作为目标序列传入
            predicted_index = torch.argmax(output, dim=-1)[0, -1].item()
            output_sequence.append(predicted_index)
            input_tensor = torch.cat([input_tensor, torch.tensor([[predicted_index]])], dim=-1)  # 将预测的单词拼接到输入序列的末尾

    # 将模型生成的单词序列转换回文本
    generated_words = [genVocab[word] for word in output_sequence if word in genVocab]
    generated_answer = ''.join(generated_words)
    return generated_answer


def digitToSentence(output_sequence):
    generated_words = [genVocab[word] for word in output_sequence.numpy() if word in genVocab]
    generated_answer = ''.join(generated_words)
    return generated_answer


num_epochs = 50
batch_size = 10
# 训练数据

news_article = "为什么女人更爱喝咖啡，而男人更爱喝酒呢？据研究表明，在一般社交场合，女人更喜欢咖啡，而男人更喜欢喝酒。在19世纪的德国，咖啡曾被讽刺为“女人的饮料”。在他们看来，女人的喋喋不休，女人的自作主张，似乎都跟咖啡有着千丝万缕的关系，虽然可爱但不能成为常态。在女人看来，咖啡是浪漫，是情调，而在男人眼里，咖啡的功能仅限于提神醒脑。而在社交场合，男人更愿意以酒示人。在他们眼里酒代表了豪迈，酒代表了阳刚之气。咖啡的口感细腻，润滑，香甜而不腻。在当今糕点跟水果似乎成了咖啡的标配，这些都是女人的首选。我们常喝的白咖啡中富含更多的蛋白质，属优质蛋白质，经由人体消化之后，分解成人体必须的氨基酸，再在水果中所含维生素C的辅助作用，合成能便皮肤细腻、光滑、富有弹性的胶原蛋白，所以说白咖啡的美容养颜的功效也便体现在此吧。这里说一下的，吃胶原蛋白能起到美容养颜的作用，是一种错误的说法。所有的蛋白质吃进去，都会经过消化分解，再经人体重新合成。要想做到美容养颜，要多食用优质蛋白质，胶原蛋白并非优质蛋白质，而水果在其中起到必不可少的辅助作用，这就是人们常说的，多吃水果能美容的原因所在。除了我们常喝的白咖啡，黑咖啡也受女士所钟爱，咖啡可以促进脂肪分解帮助燃脂，在女士减重的过程中起到一定的辅助作用。不过黑咖啡虽然对于减重有辅助作用，但是人体一天所能接受的最在摄入量，不得超过1500ml。过量摄入咖啡，会引起失眠、焦虑、烦燥；咖啡中所含的生物碱过量，会使我们的胃肠道受损。这就是说，咖啡也会醉，茶也会醉。酒是男人社交的必须品，在中国，朋友聚会喝酒，谈生意喝酒，失意喝酒，失恋喝酒等等，酒成了男人生活里不可分割的一部分，酒与中国文化密切相关。我国古代诗人似乎都与酒有着密不可分的关系，不得志喝酒，比如李白的《将进酒》便是不得志时所作；送别喝酒，比如王维的“劝君更进一杯酒，西出阳关无故人”；愁闷喝酒，比如曹操的“何以解忧，唯有杜康”，杜康是酒的发明人，这里用杜康代指酒，不过我倒认为，忧从中来，杜康何解。再比如说，李清照的“三杯两盏淡酒，怎敌他晚来风急”，李清照是我国古代著名的女词人，那个时候咖啡才刚刚被发现，还并未传到中国，如果咖啡在古代出现便传到中国，李清照的诗词里可能更多会出现咖啡也不一定哦。酒也有一定的营养价值，只不过白酒的营养价值有限，黄酒、葡萄酒的营养价值还是比较高的，适当的饮酒，对人体也是有好处的。人体所能接受的量，白酒一天摄入量不超过50ml，啤酒不超过1瓶是为最好。咖啡与酒并不冲突，酒后喝一杯咖啡，可以促进肝肾功能活跃，从而使酒精转变成乙醛快速氧化，分解成二氧化碳和水，再经由肾脏排出体外。男人爱喝酒，女人爱喝咖啡，咖啡可以解酒，酒能解忧，相辅相成，共同发展。"
word_list = split_chinese(news_article)
genVocab = {i: word for i, word in enumerate(set(word_list))}
# 创建词汇表
vocab = {word: i for i, word in enumerate(set(word_list))}
vocab_size = len(vocab)

data = []
for i in range(len(word_list) - 1):
    data.append(vocab[word_list[i]])
train_source = torch.tensor(data)
sentence = digitToSentence(train_source)
# print(sentence)
train_data = batchify(train_source, batch_size)
np_train_data = np.array(train_data)
batchSen = digitToSentence(train_data.view(-1))
# print(batchSen)
train_dataset = TextDataset(np_train_data)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=False)

# 初始化模型
model = SimpleGPT(vocab_size)

# 训练模型
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

for epoch in range(num_epochs):
    model.train()
    total_loss = 0
    for batch_idx, (input_seq, target_seq) in enumerate(train_loader):
        optimizer.zero_grad()
        # input_seq=tensor(32,9),target_seq=tensor(32,9),output=tensor(32,9,vocab_size)
        inputSen = digitToSentence(input_seq.view(-1))
        targetSen = digitToSentence(target_seq.view(-1))
        # print(inputSen)
        # print(targetSen)
        output = model(input_seq, target_seq)
        first = output.view(-1, vocab_size)
        second = target_seq.view(-1)
        # first=tensor(288,vocab_size),second=tensor(288,)
        loss = criterion(first, second)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print(f"Epoch {epoch + 1}, Average Loss: {total_loss / len(train_loader)}")


# 测试模型
# model.eval()
# test_loss = 0
# with torch.no_grad():
#     for input_seq, target_seq in test_loader:
#         output = model(input_seq, target_seq)
#         loss = criterion(output.view(-1, vocab_size), target_seq.view(-1))
#         test_loss += loss.item()
#     print(f"Test Loss: {test_loss / len(test_loader)}")

def main():
    # 创建模型和词汇表（假设已经定义）
    gpt_model = SimpleGPT(vocab_size)
    # 生成回答
    question = "人生的意义"
    answer = generate_answer(question, gpt_model, vocab)
    print("Question:", question)
    print("Answer:", answer)


if __name__ == '__main__':
    main()
    print("==========over==========")
